import cv2
import numpy as np


class color_matrix(object):
    def __init__(self):
        '''红色范围'''
        self.lower_red = None
        self.upper_red = None
        self.lower_red1 = None
        self.upper_red1 = None

        '''绿色范围'''
        self.lower_green = None
        self.upper_green = None

        '''蓝色范围'''
        self.lower_blue = None
        self.upper_blue = None
        self.frame = None

        self.Minimum_area = None  # 最小面积
        self.Maximum_angle = None  # 色块外接矩形最大旋转角度
        self.Slope = None  # 直线斜率
        self.Point_Maximum_distance = None  # 判断一个点到直线的最大距离
        self.block_x_uniformity = None  # 色块在X轴上的均与度 这个值越大容错率越大

        self.corrosion = None  # 图像腐蚀的一个系数 加大这个值可以消除图形二值化之后的杂点
        self.blurry = None
        self.left_black = None
        self.right_black = None
        self.leftmid_black = None
        self.rightmid_black = None
        self.black_wid = None

    def Program_parameters_init(self, Minimum_area, Maximum_angle, Slope, Point_Maximum_distance,
                                block_x_uniformity):  # 初始化程序参数
        self.Minimum_area = Minimum_area
        self.Maximum_angle = Maximum_angle
        self.Slope = Slope
        self.Point_Maximum_distance = Point_Maximum_distance
        self.block_x_uniformity = block_x_uniformity

    def Graphics_parameters_init(self, corrosion, blurry, left_black, right_black, leftmid_black, rightmid_black, black_wid):  # 初始化图形参数
        self.corrosion = corrosion
        self.blurry = blurry

        self.left_black = left_black
        self.right_black = right_black

        self.leftmid_black = leftmid_black
        self.rightmid_black = rightmid_black
        self.black_wid = black_wid

    def init_threshold(self, Red_threshold_low, Red_threshold_high, Red1_threshold_low, Red1_threshold_high,
                       Green_threshold_low, Green_threshold_high, Blue_threshold_low, Blue_threshold_high):
        self.lower_red = Red_threshold_low
        self.upper_red = Red_threshold_high
        self.lower_red1 = Red1_threshold_low
        self.upper_red1 = Red1_threshold_high

        self.lower_green = Green_threshold_low
        self.upper_green = Green_threshold_high

        self.lower_blue = Blue_threshold_low
        self.upper_blue = Blue_threshold_high

    def find(self, frame):

        cv2.rectangle(frame, (0, 0), (self.left_black, frame.shape[0]), (255, 255, 255), thickness=-1)
        cv2.rectangle(frame, (frame.shape[1] - self.right_black, 0), (frame.shape[1], frame.shape[0]), (255, 255, 255), thickness=-1)

        cv2.rectangle(frame, (self.leftmid_black, 0), (self.leftmid_black + self.black_wid, frame.shape[0]), (255, 255, 255), thickness=-1)
        cv2.rectangle(frame, (self.rightmid_black, 0), (self.rightmid_black + self.black_wid, frame.shape[0]), (255, 255, 255), thickness=-1)

        frame1 = frame.copy()
        frame1 = cv2.GaussianBlur(frame1, (self.blurry, self.blurry), 0)  # 高斯模糊
        hsv = cv2.cvtColor(frame1, cv2.COLOR_BGR2HSV)

        centers = []
        red_centers = []
        green_centers = []
        blue_centers = []

        mask_red1 = cv2.inRange(hsv, self.lower_red1, self.upper_red1)  # 二值化
        mask_red = cv2.inRange(hsv, self.lower_red, self.upper_red)  # 二值化
        mask_red2 = mask_red + mask_red1
        mask_red2 = cv2.erode(mask_red2, np.ones((2, 2), np.uint8), iterations=self.corrosion)
        red_contours, hierarchy = cv2.findContours(mask_red2, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

        for c in range(len(red_contours)):
            rect = cv2.minAreaRect(red_contours[c])  # 获取最小外接矩形的中心（x，y），（宽度，高度），旋转角度
            if (rect[2] < self.Maximum_angle or rect[2] > (90 - self.Maximum_angle)) and (
                    cv2.contourArea(red_contours[c]) > self.Minimum_area):
                cx, cy = rect[0]  # 中心坐标
                centers.append(rect[0])
                red_centers.append(rect[0])
                box = cv2.boxPoints(rect)  # 返回矩形顶点坐标[ [x0,y0], [x1,y1], [x2,y2], [x3,y3] ]
                box = np.int0(box)
                cv2.drawContours(frame, [box], 0, (0, 255, 0), 2)
                cv2.circle(frame, (np.int32(cx), np.int32(cy)), 2, (255, 255, 255), 2, 8, 0)

        mask_green = cv2.inRange(hsv, self.lower_green, self.upper_green)  # 二值化
        mask_green = cv2.erode(mask_green, np.ones((2, 2), np.uint8), iterations=self.corrosion)
        green_contours, hierarchy = cv2.findContours(mask_green, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

        for c in range(len(green_contours)):
            rect = cv2.minAreaRect(green_contours[c])  # 获取最小外接矩形的中心（x，y），（宽度，高度），旋转角度
            if (rect[2] < self.Maximum_angle or rect[2] > (90 - self.Maximum_angle)) and (
                    cv2.contourArea(green_contours[c]) > self.Minimum_area):
                cx, cy = rect[0]  # 中心坐标
                centers.append(rect[0])
                green_centers.append(rect[0])
                box = cv2.boxPoints(rect)  # 返回矩形顶点坐标[ [x0,y0], [x1,y1], [x2,y2], [x3,y3] ]
                box = np.int0(box)
                cv2.drawContours(frame, [box], 0, (0, 255, 0), 2)
                cv2.circle(frame, (np.int32(cx), np.int32(cy)), 2, (255, 255, 255), 2, 8, 0)

        mask_blue = cv2.inRange(hsv, self.lower_blue, self.upper_blue)  # 二值化
        mask_blue = cv2.erode(mask_blue, np.ones((2, 2), np.uint8), iterations=self.corrosion)
        blue_contours, hierarchy = cv2.findContours(mask_blue, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

        for c in range(len(blue_contours)):
            rect = cv2.minAreaRect(blue_contours[c])  # 获取最小外接矩形的中心（x，y），（宽度，高度），旋转角度
            if (rect[2] < self.Maximum_angle or rect[2] > (90 - self.Maximum_angle)) and (
                    cv2.contourArea(blue_contours[c]) > self.Minimum_area):
                cx, cy = rect[0]  # 中心坐标
                centers.append(rect[0])
                blue_centers.append(rect[0])
                box = cv2.boxPoints(rect)  # 返回矩形顶点坐标[ [x0,y0], [x1,y1], [x2,y2], [x3,y3] ]
                box = np.int0(box)
                cv2.drawContours(frame, [box], 0, (0, 255, 0), 2)
                cv2.circle(frame, (np.int32(cx), np.int32(cy)), 2, (255, 255, 255), 2, 8, 0)

        color_block = mask_blue + mask_green + mask_red2

        centers_ok1 = []
        centers_ok2 = []
        centers_ok3 = []
        centers_ok4 = []
        centers_ok5 = []
        centers_ok6 = []

        for i in range(len(centers) - 1):
            for j in range(len(centers) - i - 1):
                x_spacing = centers[i][0] - centers[j + i + 1][0]
                y_spacing = centers[i][1] - centers[j + i + 1][1]

                if x_spacing != 0:
                    k = y_spacing / x_spacing
                    if abs(k) < self.Slope:
                        centers_ok1.append([centers[i], centers[j + i + 1]])
        num = 0
        for i in centers_ok1:
            x_spacing = i[0][0] - i[1][0]
            y_spacing = i[0][1] - i[1][1]
            k = y_spacing / x_spacing
            b = i[0][1] - i[0][0] * k
            others = centers_ok1.copy()
            del others[num]
            num += 1
            for other in others:
                for one_other in other:
                    if abs(one_other[0] * k + b - one_other[1]) < self.Point_Maximum_distance and (
                            (i[0][0] > one_other[0] > i[1][0]) or (i[0][0] < one_other[0] < i[1][0])):
                        centers_ok2.append([i[0], i[1], one_other])
        for i in centers_ok2:
            if i not in centers_ok3:
                centers_ok3.append(i)
        for i in centers_ok3:
            centers_ok4.append(sorted(i, key=(lambda x: x[0])))
        for i in centers_ok4:
            if abs(abs(i[0][0] - i[1][0]) - abs(i[1][0] - i[2][0])) < self.block_x_uniformity:
                centers_ok5.append(i)
        centers_ok5 = sorted(centers_ok5, key=(lambda x: x[0][1]))
        color = ''
        if centers_ok5:
            cv2.line(frame, tuple(np.int32(centers_ok5[0][0])), tuple(np.int32(centers_ok5[0][2])), (0, 0, 255), 2)
            for i in centers_ok5:
                if i[0] in red_centers:
                    color += 'R'
                if i[0] in green_centers:
                    color += 'G'
                if i[0] in blue_centers:
                    color += 'B'

                if i[1] in red_centers:
                    color += 'R'
                if i[1] in green_centers:
                    color += 'G'
                if i[1] in blue_centers:
                    color += 'B'

                if i[2] in red_centers:
                    color += 'R'
                if i[2] in green_centers:
                    color += 'G'
                if i[2] in blue_centers:
                    color += 'B'
        return frame, color_block, color
